{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object Detection using TAO FasterRCNN\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\"> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample prediction of FasterRCNN\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/common/sample.jpg?raw=true\" width=\"960\"> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained resnet18 model and train a ResNet-18 FasterRCNN model on the KITTI dataset\n",
    "* Prune the trained FasterRCNN model\n",
    "* Retrain the pruned model to recover lost accuracy\n",
    "* Run evaluation & inference on the trained model to verify the accuracy\n",
    "* Export & deploy the model in DeepStream/TensorRT\n",
    "* Quantization-Aware Training(QAT) workflow for the best accuracy-performance trade-off\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `faster_rcnn` model\n",
    "which you may deploy via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps)\n",
    "or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    " \n",
    " ### Table of Contents\n",
    "\n",
    " This notebook shows an example use case of FasterRCNN using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    " 0. [Set up env variables and map drives](#head-0)\n",
    " 1. [Install the TAO launcher](#head-1)\n",
    " 2. [Prepare dataset and pretrained model](#head-2)<br>\n",
    "     2.1 [Download the dataset](#head-2-1)<br>\n",
    "     2.2 [Verify the downloaded dataset](#head-2-2)<br>\n",
    "     2.3 [Prepare tfrecords from kitti format dataset](#head-2-3)<br>\n",
    "     2.4 [Download pretrained model](#head-2-4)\n",
    " 3. [Provide training specification](#head-3)\n",
    " 4. [Run TAO training](#head-4)\n",
    " 5. [Evaluate trained models](#head-5)\n",
    " 6. [Prune trained models](#head-6)\n",
    " 7. [Retrain pruned models](#head-7)\n",
    " 8. [Evaluate retrained model](#head-8)\n",
    " 9. [Visualize inferences](#head-9)\n",
    " 10. [Deploy](#head-10)\n",
    " 11. [QAT workflow](#head-11)<br>\n",
    "     11.1 [Training](#head-11.1)<br>\n",
    "     11.2 [Evaluation](#head-11.2)<br>\n",
    "     11.3 [Pruning](#head-11.3)<br>\n",
    "     11.4 [Retraining](#head-11.4)<br>\n",
    "     11.5 [Evaluation of the retrained model](#head-11.5)<br>\n",
    "     11.6 [Inference of the retrained model](#head-11.6)<br>\n",
    "     11.7 [Deployment of the QAT model](#head-11.7)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    " \n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "import os\n",
    "\n",
    "print(\"Please replace the variables with your own.\")\n",
    "%env GPU_INDEX=0\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/your/tao-experiments\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"data\"\n",
    ")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"faster_rcnn\"\n",
    ")\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/faster_rcnn\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/faster_rcnn\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "%env SPECS_DIR=/workspace/tao-experiments/faster_rcnn/specs\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create local dir\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from in and out of the docker. For more information please refer to the [launcher instance](https://docs.nvidia.com/tao/tao-toolkit/text/tao_launcher.html) in the user guide.\n",
    "\n",
    "When running this cell on AWS, update the drive_map entry with the dictionary defined below, so that you don't have permission issues when writing data into folders created by the TAO docker.\n",
    "\n",
    "```json\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ],\n",
    "    \"DockerOptions\": {\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid())\n",
    "    },\n",
    "    # set gpu index for tao-converter\n",
    "    \"Envs\": [\n",
    "        {\"variable\": \"CUDA_VISIBLE_DEVICES\", \"value\": os.getenv(\"GPU_INDEX\")},\n",
    "    ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ],\n",
    "    # set gpu index for tao-converter\n",
    "    \"Envs\": [\n",
    "        {\"variable\": \"CUDA_VISIBLE_DEVICES\", \"value\": os.getenv(\"GPU_INDEX\")},\n",
    "    ]\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 460+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Skip this step if you have already installed the TAO launcher.\n",
    "!pip3 install --upgrade nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 2. Prepare dataset and pretrained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the KITTI detection dataset for the tutorial. To find more details please visit\n",
    " http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d. Please download the KITTI detection images (http://www.cvlibs.net/download.php?file=data_object_image_2.zip) and labels (http://www.cvlibs.net/download.php?file=data_object_label_2.zip) to $DATA_DOWNLOAD_DIR.\n",
    " \n",
    " The data will then be extracted to have\n",
    " * training images in `$LOCAL_DATA_DIR/training/image_2`\n",
    " * training labels in `$LOCAL_DATA_DIR/training/label_2`\n",
    " * testing images in `$LOCAL_DATA_DIR/testing/image_2`\n",
    " \n",
    "You may use this notebook with your own dataset as well. To use this example with your own dataset, please follow the same directory structure as mentioned below.\n",
    "\n",
    "*Note: There are no labels for the testing images, therefore we use it just to visualize inferences for the trained model.*"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download the dataset <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have gotten the download links in your email, please populate them in place of the `KITTI_IMAGES_DOWNLOAD_URL` and the `KITTI_LABELS_DOWNLOAD_URL`. This next cell, will download the data and place in `$LOCAL_DATA_DIR`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "os.environ[\"URL_IMAGES\"]=KITTI_IMAGES_DOWNLOAD_URL\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then wget $URL_IMAGES -O $LOCAL_DATA_DIR/data_object_image_2.zip; else echo \"image archive already downloaded\"; fi \n",
    "os.environ[\"URL_LABELS\"]=KITTI_LABELS_DOWNLOAD_URL\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then wget $URL_LABELS -O $LOCAL_DATA_DIR/data_object_label_2.zip; else echo \"label archive already downloaded\"; fi "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Verify the downloaded dataset <a class=\"anchor\" id=\"head-2-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then echo 'Image zip file not found, please download.'; else echo 'Found Image zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then echo 'Label zip file not found, please download.'; else echo 'Found Labels zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This may take a while: verify integrity of zip files \n",
    "!sha256sum $LOCAL_DATA_DIR/data_object_image_2.zip | cut -d ' ' -f 1 | grep -xq '^351c5a2aa0cd9238b50174a3a62b846bc5855da256b82a196431d60ff8d43617$' ; \\\n",
    "if test $? -eq 0; then echo \"images OK\"; else echo \"images corrupt, re-download!\" && rm -f $LOCAL_DATA_DIR/data_object_image_2.zip; fi \n",
    "!sha256sum $LOCAL_DATA_DIR/data_object_label_2.zip | cut -d ' ' -f 1 | grep -xq '^4efc76220d867e1c31bb980bbf8cbc02599f02a9cb4350effa98dbb04aaed880$' ; \\\n",
    "if test $? -eq 0; then echo \"labels OK\"; else echo \"labels corrupt, re-download!\" && rm -f $LOCAL_DATA_DIR/data_object_label_2.zip; fi "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack \n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_image_2.zip -d $LOCAL_DATA_DIR\n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_label_2.zip -d $LOCAL_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "import os\n",
    "\n",
    "DATA_DIR = os.environ.get('LOCAL_DATA_DIR')\n",
    "num_training_images = len(os.listdir(os.path.join(DATA_DIR, \"training/image_2\")))\n",
    "num_training_labels = len(os.listdir(os.path.join(DATA_DIR, \"training/label_2\")))\n",
    "num_testing_images = len(os.listdir(os.path.join(DATA_DIR, \"testing/image_2\")))\n",
    "print(\"Number of images in the train/val set. {}\".format(num_training_images))\n",
    "print(\"Number of labels in the train/val set. {}\".format(num_training_labels))\n",
    "print(\"Number of images in the test set. {}\".format(num_testing_images))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Prepare tfrecords from kitti format dataset <a class=\"anchor\" id=\"head-2-3\"></a>\n",
    "\n",
    "* Update the tfrecords spec file to take in your kitti format dataset\n",
    "* Create the tfrecords using the dataset_convert \n",
    "* TFRecords only need to be generated once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TFrecords conversion spec file for training\")\n",
    "!cat $LOCAL_SPECS_DIR/frcnn_tfrecords_kitti_trainval.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#KITTI trainval\n",
    "!tao model faster_rcnn dataset_convert --gpu_index $GPU_INDEX -d $SPECS_DIR/frcnn_tfrecords_kitti_trainval.txt \\\n",
    "                     -o $DATA_DOWNLOAD_DIR/faster_rcnn/tfrecords/kitti_trainval/kitti_trainval \\\n",
    "                     -r $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_DATA_DIR/faster_rcnn/tfrecords/kitti_trainval"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 2.4 Download pre-trained model <a class=\"anchor\" id=\"head-2-4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/pretrained_object_detection*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download model from NGC.\n",
    "!ngc registry model download-version nvidia/tao/pretrained_object_detection:resnet18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Copy weights to experiment directory.\n",
    "!cp pretrained_object_detection_vresnet18/resnet_18.hdf5 $LOCAL_EXPERIMENT_DIR\n",
    "!rm -rf pretrained_object_detection_vresnet18\n",
    "!ls -rlt $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/default_spec_resnet18.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    " * Provide the sample spec file for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model faster_rcnn train --gpu_index $GPU_INDEX -e $SPECS_DIR/default_spec_resnet18.txt -r /workspace/tao-experiments/faster_rcnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model for each epoch:')\n",
    "print('---------------------')\n",
    "!ls -lht $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"For multi-GPU data parallelism, please uncomment and run this instead. Change --gpus  and --gpu_index based on your machine.\")\n",
    "# !tao model faster_rcnn train -e $SPECS_DIR/default_spec_resnet18.txt \\\n",
    "#                    --gpus 2 \\\n",
    "#                    --gpu_index 0 1 \\\n",
    "#                    -r /workspace/tao-experiments/faster_rcnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\"\"\n",
    "For multi-GPU model parallelism, please uncomment and run this instead.\n",
    "Also add related parameters in training_config to enable model parallelism. E.g., \n",
    "\n",
    "             model_parallelism: 50\n",
    "             model_parallelism: 50\n",
    "\n",
    "\"\"\")\n",
    "\n",
    "#!tao model faster_rcnn train -e $SPECS_DIR/default_spec_resnet18.txt \\\n",
    "#                   --gpus 2 \\\n",
    "#                   --gpu_index 0 1\\\n",
    "#                   -np 1 \\\n",
    "#                   -r /workspace/tao-experiments/faster_rcnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"For resume training from checkpoint, please uncomment and run this instead. Change/Add the 'resume_from_model' field in the spec file.\")\n",
    "# !tao model faster_rcnn train --gpu_index $GPU_INDEX -e $SPECS_DIR/default_spec_resnet18.txt -r /workspace/tao-experiments/faster_rcnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"For Automatic Mixed Precision(AMP) training, please uncomment and run this. Make sure you use the Volta or above GPU arch to enable AMP.\")\n",
    "# !tao model faster_rcnn train --gpu_index $GPU_INDEX --use_amp -e $SPECS_DIR/default_spec_resnet18.txt -r /workspace/tao-experiments/faster_rcnn"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model faster_rcnn evaluate --gpu_index $GPU_INDEX -e $SPECS_DIR/default_spec_resnet18.txt -m /workspace/tao-experiments/faster_rcnn/frcnn_resnet_18.epoch_12.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 6. Prune trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    " * Specify pre-trained model\n",
    " * Equalization criterion\n",
    " * Threshold for pruning\n",
    " * Output directory to store the model\n",
    " \n",
    "Usually, you just need to adjust `-pth` (threshold) for accuracy and model size trade off. Higher `pth` gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold to use is depend on the dataset. A `pth` value below is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model faster_rcnn prune --gpu_index $GPU_INDEX -m $USER_EXPERIMENT_DIR/frcnn_resnet_18.epoch_12.hdf5 \\\n",
    "           -o $USER_EXPERIMENT_DIR/model_1_pruned.hdf5  \\\n",
    "           -eq union  \\\n",
    "           -pth 0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -lht $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 7. Retrain pruned models <a class=\"anchor\" id=\"head-7\"></a>\n",
    " * Model needs to be re-trained to bring back accuracy after pruning\n",
    " * Specify re-training specification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we have updated the spec file to include the newly pruned model as a pretrained weights.\n",
    "!cat $LOCAL_SPECS_DIR/default_spec_resnet18_retrain_spec.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights \n",
    "!tao model faster_rcnn train --gpu_index $GPU_INDEX -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt -r /workspace/tao-experiments/faster_rcnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -lht $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 8. Evaluate retrained model <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model faster_rcnn evaluate --gpu_index $GPU_INDEX -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt -m /workspace/tao-experiments/faster_rcnn/frcnn_resnet_18.epoch_12.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 9. Visualize inferences <a class=\"anchor\" id=\"head-9\"></a>\n",
    " In this section, we run the inference tool to generate inferences on the trained models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy some test images\n",
    "!mkdir -p $LOCAL_DATA_DIR/test_samples\n",
    "!cp $LOCAL_DATA_DIR/testing/image_2/00000* $LOCAL_DATA_DIR/test_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference for detection on n images\n",
    "# Please go to $LOCAL_EXPERIMENT_DIR/inference_results_imgs_retrain to see the visualizations.\n",
    "!tao model faster_rcnn inference --gpu_index $GPU_INDEX \\\n",
    "                -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                -m /workspace/tao-experiments/faster_rcnn/frcnn_resnet_18.epoch_12.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `inference` tool produces two outputs. \n",
    "1. Overlain images in `$LOCAL_EXPERIMENT_DIR/inference_results_imgs_retrain`\n",
    "2. Frame by frame bbox labels in kitti format located in `$LOCAL_EXPERIMENT_DIR/inference_dump_labels_retrain`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg', '.png', '.jpeg', '.ppm']\n",
    "\n",
    "def visualize_images(image_dir, num_cols=4, num_images=10):\n",
    "    output_path = os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'], image_dir)\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'inference_results_imgs_retrain' # relative path from $LOCAL_EXPERIMENT_DIR.\n",
    "COLS = 3 # number of columns in the visualizer grid.\n",
    "IMAGES = 9 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 10. Deploy! <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Generate .onnx file using tao container\n",
    "!if [ -f $LOCAL_EXPERIMENT_DIR/frcnn_resnet_18.onnx ]; then rm -f $LOCAL_EXPERIMENT_DIR/frcnn_resnet_18.onnx; fi\n",
    "!tao model faster_rcnn export --gpu_index $GPU_INDEX -m $USER_EXPERIMENT_DIR/frcnn_resnet_18.epoch_12.hdf5  \\\n",
    "                        -o $USER_EXPERIMENT_DIR/frcnn_resnet_18.onnx \\\n",
    "                        -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                        --target_opset 12 \\\n",
    "                        --gen_ds_config"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `tao deploy` container, you can generate a TensorRT engine and verify the correctness of the generated through evaluate and inference. \n",
    "\n",
    "The `tao deploy` produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please run `tao deploy` command which will instantiate a deploy container, with the exported `.onnx` file on your target device. The `tao deploy` container only works for x86, with discrete NVIDIA GPU's. \n",
    "\n",
    "For the jetson devices, please download the tao-converter for jetson from the dev zone link [here](https://developer.nvidia.com/tao-converter). \n",
    "\n",
    "If you choose to integrate your model into deepstream directly, you may do so by simply copying the exported `.onnx` file along with the calibration cache to the target device and updating the spec file that configures the `gst-nvinfer` element to point to this newly exported model. Usually this file is called `config_infer_primary.txt` for detection models and `config_infer_secondary_*.txt` for classification models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (FP32).\n",
    "!tao deploy faster_rcnn gen_trt_engine --gpu_index $GPU_INDEX \\\n",
    "                        -m $USER_EXPERIMENT_DIR/frcnn_resnet_18.onnx \\\n",
    "                        -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                        --data_type fp32 \\\n",
    "                        --batch_size 8 \\\n",
    "                        --max_batch_size 4 \\\n",
    "                        --engine_file $USER_EXPERIMENT_DIR/trt.fp32.engine \\\n",
    "                        --results_dir $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (FP16).\n",
    "!tao deploy faster_rcnn gen_trt_engine --gpu_index $GPU_INDEX \\\n",
    "                        -m $USER_EXPERIMENT_DIR/frcnn_resnet_18.onnx \\\n",
    "                        -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                        --data_type fp16 \\\n",
    "                        --batch_size 8 \\\n",
    "                        --max_batch_size 4 \\\n",
    "                        --engine_file $USER_EXPERIMENT_DIR/trt.fp16.engine \\\n",
    "                        --results_dir $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (INT8).\n",
    "!tao deploy faster_rcnn gen_trt_engine --gpu_index $GPU_INDEX \\\n",
    "                        -m $USER_EXPERIMENT_DIR/frcnn_resnet_18.onnx \\\n",
    "                        -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                        --data_type int8 \\\n",
    "                        --batch_size 8 \\\n",
    "                        --max_batch_size 4 \\\n",
    "                        --batches 10 \\\n",
    "                        --cal_cache_file $USER_EXPERIMENT_DIR/cal.bin \\\n",
    "                        --cal_image_dir $DATA_DOWNLOAD_DIR/testing/image_2 \\\n",
    "                        --cal_data_file $USER_EXPERIMENT_DIR/cal.tensorfile \\\n",
    "                        --engine_file $USER_EXPERIMENT_DIR/trt.int8.engine \\\n",
    "                        --results_dir $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model and converted TensorRT engine:')\n",
    "print('------------')\n",
    "!ls -lht $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do inference with TensorRT on the generated TensorRT engine\n",
    "# Please go to $LOCAL_EXPERIMENT_DIR/images_annotated to see the visualizations.\n",
    "!tao deploy faster_rcnn inference  --gpu_index $GPU_INDEX \\\n",
    "                                   -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                                   -m $USER_EXPERIMENT_DIR/trt.fp32.engine \\\n",
    "                                   -i $DATA_DOWNLOAD_DIR/test_samples \\\n",
    "                                   -r $USER_EXPERIMENT_DIR/trt_infer"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `inference` tool produces two outputs. \n",
    "The paths to the two outputs are exactly the same as the first `inference` command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images from TensorRT inference.\n",
    "OUTPUT_PATH = 'inference_results_imgs_retrain' # relative path from $LOCAL_EXPERIMENT_DIR.\n",
    "COLS = 3 # number of columns in the visualizer grid.\n",
    "IMAGES = 9 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "val_split = 0.14  # Val split from the spec file\n",
    "img_ext = '.png'\n",
    "lab_ext = '.txt'\n",
    "\n",
    "# Get images and labels from the partitioned validation set\n",
    "images_root = f\"{os.environ['LOCAL_DATA_DIR']}/training/image_2\"\n",
    "labels_root = f\"{os.environ['LOCAL_DATA_DIR']}/training/label_2\"\n",
    "images_list = [os.path.splitext(imfile)[0] for imfile in sorted(os.listdir(images_root)) if imfile.endswith(img_ext)]\n",
    "num_val_images = int(len(images_list) * val_split)\n",
    "\n",
    "# Copy the data to a separate directory for evaluation\n",
    "os.makedirs(os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"images_val\"), exist_ok=True)\n",
    "os.makedirs(os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"labels_val\"), exist_ok=True)\n",
    "\n",
    "for fname in images_list[:num_val_images]:\n",
    "    shutil.copyfile(os.path.join(images_root, fname + img_ext), \\\n",
    "                             os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"images_val\", fname + img_ext))\n",
    "    shutil.copyfile(os.path.join(labels_root, fname + lab_ext), \\\n",
    "                             os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"labels_val\", fname + lab_ext))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Doing evaluation with the generated TensorRT engine\n",
    "!tao deploy faster_rcnn evaluate --gpu_index $GPU_INDEX \\\n",
    "                                 -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                                 -m $USER_EXPERIMENT_DIR/trt.int8.engine \\\n",
    "                                 -i $DATA_DOWNLOAD_DIR/training/images_val \\\n",
    "                                 -l $DATA_DOWNLOAD_DIR/training/labels_val \\\n",
    "                                 -r $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 11. QAT workflow <a class=\"anchor\" id=\"head-11\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we will explore the typical Quantization-Aware Training(QAT) workflow with TAO. QAT workflow is almost the same as non-QAT workflow except for two major differences:\n",
    "1. set `enable_qat` to `True` in training and retraining spec files to enable the QAT for training/retraining\n",
    "2. when doing export in INT8 mode, the calibration cache is extracted directly from the QAT .hdf5 model, so no need to specify any TensorRT INT8 calibration related arguments for `export`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 11.1. Training <a class=\"anchor\" id=\"head-10.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set enable_qat to True in training spec file to enable QAT training\n",
    "!sed -i 's/enable_qat: False/enable_qat: True/' $LOCAL_SPECS_DIR/default_spec_resnet18.txt\n",
    "!cat $LOCAL_SPECS_DIR/default_spec_resnet18.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run QAT training\n",
    "!tao model faster_rcnn train --gpu_index $GPU_INDEX -e $SPECS_DIR/default_spec_resnet18.txt -r /workspace/tao-experiments/faster_rcnn"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 11.2. Evaluation <a class=\"anchor\" id=\"head-10.2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model faster_rcnn evaluate --gpu_index $GPU_INDEX -e $SPECS_DIR/default_spec_resnet18.txt -m /workspace/tao-experiments/faster_rcnn/frcnn_resnet_18.epoch_12.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 11.3. Pruning <a class=\"anchor\" id=\"head-10.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model faster_rcnn prune --gpu_index $GPU_INDEX -m $USER_EXPERIMENT_DIR/frcnn_resnet_18.epoch_12.hdf5 \\\n",
    "           -o $USER_EXPERIMENT_DIR/model_1_pruned.hdf5  \\\n",
    "           -eq union  \\\n",
    "           -pth 0.2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 11.4. Retraining <a class=\"anchor\" id=\"head-10.4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set enable_qat to True in retraining spec file to enable QAT\n",
    "!sed -i 's/enable_qat: False/enable_qat: True/' $LOCAL_SPECS_DIR/default_spec_resnet18_retrain_spec.txt\n",
    "!cat $LOCAL_SPECS_DIR/default_spec_resnet18_retrain_spec.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model faster_rcnn train --gpu_index $GPU_INDEX -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt -r /workspace/tao-experiments/faster_rcnn"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 11.5. Evaluation of the retrained model <a class=\"anchor\" id=\"head-10.5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# do evaluation with .hdf5 model\n",
    "!tao model faster_rcnn evaluate --gpu_index $GPU_INDEX -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt -m /workspace/tao-experiments/faster_rcnn/frcnn_resnet_18.epoch_12.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 11.6. Inference of the retrained model <a class=\"anchor\" id=\"head-10.6\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# do inference with .hdf5 model\n",
    "!tao model faster_rcnn inference --gpu_index $GPU_INDEX \\\n",
    "                -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                -m /workspace/tao-experiments/faster_rcnn/frcnn_resnet_18.epoch_12.hdf5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images\n",
    "OUTPUT_PATH = 'inference_results_imgs_retrain' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 3 # number of columns in the visualizer grid.\n",
    "IMAGES = 9 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 11.7. Deployment of the QAT model <a class=\"anchor\" id=\"head-10.7\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calibration JSON file is required for INT8 engine generation from tao deploy\n",
    "!if [ -f $LOCAL_EXPERIMENT_DIR/frcnn_resnet_18_int8_qat.onnx ]; then rm -f $LOCAL_EXPERIMENT_DIR/frcnn_resnet_18_int8_qat.onnx; fi\n",
    "!if [ -f $LOCAL_EXPERIMENT_DIR/cal.bin ]; then rm -f $LOCAL_EXPERIMENT_DIR/cal.bin; fi\n",
    "!tao model faster_rcnn export --gpu_index $GPU_INDEX \\\n",
    "                        -m $USER_EXPERIMENT_DIR/frcnn_resnet_18.epoch_12.hdf5  \\\n",
    "                        -o $USER_EXPERIMENT_DIR/frcnn_resnet_18_int8_qat.onnx \\\n",
    "                        -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                        --cal_json_file $USER_EXPERIMENT_DIR/cal.json \\\n",
    "                        --target_opset 12 \\\n",
    "                        --gen_ds_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine(INT8).\n",
    "# No need for calibration dataset for QAT model INT8 export\n",
    "!tao deploy faster_rcnn gen_trt_engine --gpu_index $GPU_INDEX \\\n",
    "                        -m $USER_EXPERIMENT_DIR/frcnn_resnet_18_int8_qat.onnx \\\n",
    "                        -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                        --data_type int8 \\\n",
    "                        --cal_cache_file $USER_EXPERIMENT_DIR/cal.bin \\\n",
    "                        --cal_json_file $USER_EXPERIMENT_DIR/cal.json \\\n",
    "                        --batch_size 8 \\\n",
    "                        --max_batch_size 4 \\\n",
    "                        --engine_file $USER_EXPERIMENT_DIR/trt.int8.qat.engine \\\n",
    "                        --results_dir $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model and converted TensorRT engine:')\n",
    "print('------------')\n",
    "!ls -lht $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do inference with TensorRT on the generated TensorRT engine\n",
    "# Please go to $LOCAL_EXPERIMENT_DIR/images_annotated to see the visualizations.\n",
    "!tao deploy faster_rcnn inference --gpu_index $GPU_INDEX \\\n",
    "                                  -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                                  -m $USER_EXPERIMENT_DIR/trt.int8.qat.engine \\\n",
    "                                  -i $DATA_DOWNLOAD_DIR/test_samples \\\n",
    "                                  -r $USER_EXPERIMENT_DIR/trt_qat_infer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images from TensorRT inference.\n",
    "OUTPUT_PATH = 'inference_results_imgs_retrain' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 3 # number of columns in the visualizer grid.\n",
    "IMAGES = 9 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Doing evaluation with the generated TensorRT engine\n",
    "!tao deploy faster_rcnn evaluate --gpu_index $GPU_INDEX \\\n",
    "                                 -e $SPECS_DIR/default_spec_resnet18_retrain_spec.txt \\\n",
    "                                 -m $USER_EXPERIMENT_DIR/trt.int8.qat.engine \\\n",
    "                                 -i $DATA_DOWNLOAD_DIR/training/images_val \\\n",
    "                                 -l $DATA_DOWNLOAD_DIR/training/labels_val \\\n",
    "                                 -r $USER_EXPERIMENT_DIR/"
   ]
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython2",
  "version": 2,
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
